package com.kingsoft.dc.khaos.plugin.core.utils;


import com.kingsoft.dc.khaos.plugin.core.constants.Contants;

import com.kingsoft.dc.khaos.plugin.core.constants.Contants.ResourceStrategy;
import java.util.HashMap;
import java.util.Map;
import java.util.regex.Pattern;

/**
 * 资源设置处理类.
 */
public class ResourceUtils {
    static Map<String, Long> memUnitMapping = new HashMap<>();
    public static String GB = "GB";
    public static String MB = "MB";
    public static String KB = "KB";
    public static String B = "B";

    static {
        memUnitMapping.put(GB, 1024 * 1024 * 1024L);
        memUnitMapping.put(MB, 1024 * 1024L);
        memUnitMapping.put(KB, 1024L);
        memUnitMapping.put(B, 1L);
    }

    /**
     * 根据CU数量获取executor的数量
     * @param cuNums
     * @return
     */
    public static int getExecutorNum(int cuNums) {
//        int executors = Math.max(cuNums - 1 , 2);
        int executors = cuNums;
        int ratio = Integer.parseInt(ConfigUtils.pro.getProperty(Contants.EXECUTOR_CU_RATIO, "1"));
        return executors/ratio;
    }

    /**
     * 获取executor的最小数量
     * @return
     */
    public static int getMinExecutorNum() {
        return Integer.parseInt(ConfigUtils.pro.getProperty(Contants.MIN_EXECUTOR_NUMBER, "1"));
    }


    /**
     * 获取每个executor的cores
     * @return
     */
    public static int getNumCoresPerExecutor() {
        int coresPerCU = getNumCoresPerCU();
        int ratio = Integer.parseInt(ConfigUtils.pro.getProperty(Contants.EXECUTOR_CU_RATIO, "1"));
        return coresPerCU * ratio;
    }

    /**
     * 获取每个executor的cores
     * @param executorCuNum 执行节点的cu数
     * @param resourceStrategy
     * @return
     */
    public static int getNumCoresPerExecutor(int executorCuNum, String resourceStrategy) {
        int coresPerCU = getNumCoresPerCU();
        int ratio = Integer.parseInt(ConfigUtils.pro.getProperty(Contants.EXECUTOR_CU_RATIO, "1"));

        int cores =
            ResourceStrategy.CPU.name().equalsIgnoreCase(resourceStrategy) ? (coresPerCU * executorCuNum
                * ratio) : (coresPerCU * ratio);

        return cores * ratio;
    }

    /**
     * 获取每个executor的内存
     * @return
     */
    public static String getMemPerExecutorWithUnit() {
        // 资源分配模式，默认CPU优先
        return getMemPerExecutorWithUnit(1, ResourceStrategy.CPU.name());
    }

    /**
     * 获取每个executor的内存 executorCuNum,resourceStrategy
     * cpu优先：executor-mem=memPerCU(即3G）  内存优先：memPerCU(即3G） * executorCuNum
     * @return
     */
    public static String getMemPerExecutorWithUnit(int executorCuNum, String resourceStrategy) {
        long memPerCU = getMemPerCU();
        int ratio = Integer.parseInt(ConfigUtils.pro.getProperty(Contants.EXECUTOR_CU_RATIO, "1"));

        // CPU优先模式或者内存优先模式 TODO xixi 考虑堆外内存
        //long memSize = memPerCU * cuNums * ratio;
        long memSize = ResourceStrategy.CPU.name().equalsIgnoreCase(resourceStrategy) ? (memPerCU * ratio) : (memPerCU * executorCuNum * ratio);
        return memSize / memUnitMapping.get(MB) + MB;
    }

    /**
     * 获取Driver的内存
     * @return
     */
    public static int getDriverCores() {
        int memPerCU = getNumCoresPerCU();
        return memPerCU;
    }

    /**
     * 获取Driver的内存
     * @return
     */
    public static String getDriverMemory() {
        return getDriverMemory(1);
    }

    /**
     * 获取Driver的内存
     *
     * @return
     */
    public static String getDriverMemory(int cuNums) {
        long memPerCU = getMemPerCU();
        return (memPerCU * cuNums) / memUnitMapping.get(MB) + MB;
    }


    public static Long getMemPerCU() {
        String memWithUnit = ConfigUtils.pro.getProperty(Contants.MEMORY_PER_CU, "2GB").trim();
        return getMemoryBytesPerCU(memWithUnit);
    }

    /**
     * 获取每个CU的cores
     * @return
     */
    public static int getNumCoresPerCU() {
        int coresPerCU = Integer.parseInt(ConfigUtils.pro.getProperty(Contants.CORES_PER_CU, "1"));
        return coresPerCU;
    }
    public static Long getMemoryBytesPerCU (String memWithUnit) {
        long mem = 0;
        String unit = "";
        if (isInteger(memWithUnit)) {
            mem = Long.valueOf(memWithUnit);
        } else {
            unit = memWithUnit.replaceAll("([1-9]+[0-9]*|0)(\\.[\\d]+)?", "");
            String memStr = extractNumber(memWithUnit);
            if (memUnitMapping.containsKey(unit)) {
                mem = Long.valueOf(memUnitMapping.get(unit) * Integer.parseInt(memStr));
            }
        }
        return mem;
    }
    /** 判断是否为整数
    * @param str 传入的字符串
    * @return 是整数返回true,否则返回false
    */
    public static boolean isInteger(String str) {
        Pattern pattern = Pattern.compile("^[-\\+]?[\\d]*$");
        return pattern.matcher(str).matches();
    }


    /** 判断是否为整数
     * @param str 传入的字符串
     * @return 是整数返回true,否则返回false
     */
    public static String extractNumber(String str) {
        Pattern pattern = Pattern.compile("[^0-9]");
        return pattern.matcher(str).replaceAll("").trim();
    }


    public static void  main(String[] args) {
        String mem = getMemPerExecutorWithUnit();
    }

}
